/*
 * DirichletFit.cpp
 *
 *  Created on: Apr 5, 2011
 *      Author: clint
 */

#include "DirichletFit.h"

DirichletFit::DirichletFit(mat data) {
	this->dirichlet_samples_ = data;
}

DirichletFit::~DirichletFit() {}


/*
 * Function  :  dirichlet_moment_match_row(mat data)
 * Reference : 	dirichlet_moment_match.m
 *
 * Each row of data is a multivariate
 * observation on the probability simplex
 *
 */
rowvec DirichletFit::dirichlet_moment_match_row(){

	rowvec mu_alpha_vec = mean(this->dirichlet_samples_, 0);
	rowvec mu2_alpha_vec = mean(this->dirichlet_samples_ % this->dirichlet_samples_, 0);
	rowvec s = zeros<rowvec>(this->dirichlet_samples_.n_rows);

	for(size_t i = 0; i < this->dirichlet_samples_.n_rows; i++)
		if((mu2_alpha_vec(i) - mu_alpha_vec(i) * mu_alpha_vec(i)) > 0)
			s(i) = (mu_alpha_vec(i) - mu2_alpha_vec(i)) / (mu2_alpha_vec(i) - mu_alpha_vec(i) * mu_alpha_vec(i));

	return mu_alpha_vec * median(s);
}

/*
 * Function  :  dirichlet_moment_match_row(mat data)
 * Reference : 	dirichlet_moment_match.m
 *
 * Each column of data is a multivariate
 * observation on the probability simplex.
 *
 */
vec DirichletFit::dirichlet_moment_match_col(){ //

	vec mu_alpha_vec = mean(this->dirichlet_samples_, 1);
	vec mu2_alpha_vec = mean(this->dirichlet_samples_ % this->dirichlet_samples_, 1);
	vec s = zeros<vec>(this->dirichlet_samples_.n_rows);

	for(size_t i = 0; i < this->dirichlet_samples_.n_rows; i++)
		if((mu2_alpha_vec(i) - mu_alpha_vec(i) * mu_alpha_vec(i)) > 0)
			s(i) = (mu_alpha_vec(i) - mu2_alpha_vec(i)) / (mu2_alpha_vec(i) - mu_alpha_vec(i) * mu_alpha_vec(i));


	return mu_alpha_vec * median(s);
}

/*
 * Function  :  dirichlet_log_likelihood_row(rowvec alpha_vec, rowvec ln_bar_p)
 * Reference : 	dirichlet_logProb_fast.m
 *
 */
double DirichletFit::dirichlet_log_likelihood_row(rowvec alpha_vec, rowvec ln_bar_p){

	double sum_lg = 0.0;
	for (size_t i = 0; i < alpha_vec.n_elem; i++)
		sum_lg += gsl_sf_lngamma(alpha_vec(i));

	return gsl_sf_lngamma(accu(alpha_vec)) - sum_lg + accu((alpha_vec - 1) % ln_bar_p);

}

double DirichletFit::dirichlet_log_likelihood_col(vec alpha_vec, vec ln_bar_p){

	double sum_lg = 0.0;
	for (size_t i = 0; i < alpha_vec.n_elem; i++)
		sum_lg += gsl_sf_lngamma(alpha_vec(i));

	return gsl_sf_lngamma(accu(alpha_vec)) - sum_lg + accu((alpha_vec - 1) % ln_bar_p);
}


/*
 * Reference: see Eq. 12 to Eq. 18
 *
 */
rowvec DirichletFit::hessian_times_gradient_row(rowvec alpha_vec, rowvec gradient_ll, double lambda){

	double alpha_sum = accu(alpha_vec);
	double z = trigamma(alpha_sum) + 1e-24;

	rowvec q_inv_vec = zeros<rowvec>(alpha_vec.n_elem);
	for (size_t i = 0; i < alpha_vec.n_elem; i++)
		q_inv_vec(i) = 1 / (-trigamma(alpha_vec(i)) - lambda);

	double b = accu(gradient_ll % q_inv_vec) / (1/z + accu(q_inv_vec));

	return (gradient_ll - b) % q_inv_vec;

}

vec DirichletFit::hessian_times_gradient_col(vec alpha_vec, vec gradient_ll, double lambda){

	double alpha_sum = accu(alpha_vec);
	double z = trigamma(alpha_sum) + 1e-24;
	vec q_inv_vec = zeros<vec>(alpha_vec.n_elem);

	for (size_t i = 0; i < alpha_vec.n_elem; i++)
		q_inv_vec(i) = 1 / (-trigamma(alpha_vec(i)) - lambda);

	double b = accu(gradient_ll % q_inv_vec) / (1/z + accu(q_inv_vec));

	return (gradient_ll - b) % q_inv_vec;

}

bool DirichletFit::rowvec_less_than(rowvec a, rowvec b){
	for (size_t i = 0; i < a.n_elem; i++)
		if (a(i) > b(i)) return false;
	return true;
}

bool DirichletFit::vec_less_than(vec a, vec b){
	for (size_t i = 0; i < a.n_elem; i++)
		if (a(i) > b(i)) return false;
	return true;
}

/*
 * Function    : 	dirichlet_fit_newton()
 *
 * Description : 	Maximum Likelihood Dirichlet distribution based on the approach
 * 					described in "Estimating a Dirichlet distribution" by Tom Minka.
 *
 * Reference   :	fastfit/dirichlet_fit_newton.m
 *
 */

/*
 * Each row of data is a multivariate
 * observation on the probability simplex.
 */

rowvec DirichletFit::dirichlet_fit_newton_row(rowvec alpha_vec){

	rowvec mean_ln_p = mean(log(this->dirichlet_samples_), 0); // mean of each column
	size_t num_samples = this->dirichlet_samples_.n_rows;
	size_t num_dimensions = this->dirichlet_samples_.n_cols;
	double old_energy = num_samples * dirichlet_log_likelihood_row(alpha_vec, mean_ln_p);
	double new_energy = 0;
	double lambda = 0.1;
	rowvec old_alpha_vec;
	rowvec hg;
	rowvec gradient_ll = zeros<rowvec>(num_dimensions);

	for (size_t iter = 0; iter < 100; iter++){

		if (accu(alpha_vec) == 0) break;
		old_alpha_vec = alpha_vec;

		// Calculates the gradient of the log-likelihood, Eq. 6
		double psi_sum = digamma(accu(alpha_vec));
		for (size_t i = 0; i < num_dimensions; i++)
			gradient_ll(i) = psi_sum - digamma(alpha_vec(i)) + mean_ln_p(i);

		bool abort = false;

		while (1){
			hg = hessian_times_gradient_row(alpha_vec, gradient_ll, lambda);

			if (rowvec_less_than(hg, alpha_vec)){
				new_energy = num_samples * dirichlet_log_likelihood_row((alpha_vec - hg), mean_ln_p);
				if (new_energy > old_energy) {
					old_energy = new_energy;
					alpha_vec = alpha_vec - hg;
					lambda /= 10;
					break;
				}
			}

			lambda *= 10;

			if (lambda > CONST_LAMBDA_MAX) {
				abort = true;
				break;
			}
		}

		if (abort) break;

		double max_abs = 0.0;
		for (size_t i = 0; i < num_dimensions; i++){
			if (alpha_vec(i) < CONST_EPS)
				alpha_vec(i) = CONST_EPS;
			if (max_abs < abs(alpha_vec(i) - old_alpha_vec(i)))
				max_abs = abs(alpha_vec(i) - old_alpha_vec(i));
		}
		if (max_abs < CONST_ALPHA_MIN) break;
	}

	return alpha_vec;

}

rowvec DirichletFit::dirichlet_fit_newton_row(){
	rowvec alpha_vec = dirichlet_moment_match_row();
	return dirichlet_fit_newton_row(alpha_vec);
}


/*
 * Each row of data is a multivariate
 * observation on the probability simplex.
 */

vec DirichletFit::dirichlet_fit_newton_col(vec alpha_vec){

	vec mean_ln_p = mean(log(this->dirichlet_samples_), 1);
	double num_samples = this->dirichlet_samples_.n_cols;
	size_t num_dimensions = this->dirichlet_samples_.n_rows;
	double old_energy = num_samples * dirichlet_log_likelihood_col(alpha_vec, mean_ln_p);
	double new_energy = 0;
	double lambda = 0.1;
	vec old_alpha_vec;
	vec hg;
	vec gradient_ll = zeros<vec>(num_dimensions);

	for (size_t iter = 0; iter < 100; iter++){

		if (accu(alpha_vec) == 0) break;
		old_alpha_vec = alpha_vec;

		// Calculates the gradient of the log-likelihood, Eq. 6
		double psi_sum = digamma(accu(alpha_vec));
		for (size_t i = 0; i < num_dimensions; i++)
			gradient_ll(i) = psi_sum - digamma(alpha_vec(i)) + mean_ln_p(i);

		bool abort = false;

		while (1){
			hg = hessian_times_gradient_col(alpha_vec, gradient_ll, lambda);

			if (vec_less_than(hg, alpha_vec)){
				new_energy = num_samples * dirichlet_log_likelihood_col((alpha_vec - hg), mean_ln_p);
				if (new_energy > old_energy) {
					old_energy = new_energy;
					alpha_vec = alpha_vec - hg;
					lambda /= 10;
					break;
				}
			}

			lambda *= 10;

			if (lambda > CONST_LAMBDA_MAX) {
				abort = true;
				break;
			}
		}

		if (abort) break;

		double max_abs = 0.0;
		for (size_t i = 0; i < num_dimensions; i++){
			if (alpha_vec(i) < CONST_EPS)
				alpha_vec(i) = CONST_EPS;
			if (max_abs < abs(alpha_vec(i) - old_alpha_vec(i)))
				max_abs = abs(alpha_vec(i) - old_alpha_vec(i));
		}
		if (max_abs < CONST_ALPHA_MIN) break;
	}

	return alpha_vec;

}

vec DirichletFit::dirichlet_fit_newton_col(){
	vec alpha_vec = dirichlet_moment_match_col();
	return dirichlet_fit_newton_col(alpha_vec);
}


//****************************************************************************80
//
//  Purpose:
//
//    TRIGAMMA calculates trigamma(x) = d**2 log(gamma(x)) / dx**2
//
//  Modified:
//
//    19 January 2008
//
//  Author:
//
//    Original FORTRAN77 version by BE Schneider.
//    C++ version by John Burkardt.
//
//  Reference:
//
//    BE Schneider,
//    Algorithm AS 121:
//    Trigamma Function,
//    Applied Statistics,
//    Volume 27, Number 1, pages 97-99, 1978.
//
//  Parameters:
//
//    Input, double X, the argument of the trigamma function.
//    0 < X.
//
//    Output, int *IFAULT, error flag.
//    0, no error.
//    1, X <= 0.
//
//    Output, double TRIGAMMA, the value of the trigamma function at X.
//

double DirichletFit::trigamma(double x) {

	double a = 0.0001;
	double b = 5.0;
	double b2 = 0.1666666667;
	double b4 = -0.03333333333;
	double b6 = 0.02380952381;
	double b8 = -0.03333333333;
	double value;
	double y;
	double z = x;
	//
	//  Check the input.
	//
	if (x <= 0.0) return 0.0;

	z = x;
	//
	//  Use small value approximation if X <= A.
	//
	if (x <= a) {
		value = 1.0 / x / x;
		return value;
	}
	//
	//  Increase argument to ( X + I ) >= B.
	//
	value = 0.0;

	while (z < b) {
		value += 1.0 / z / z;
		z += 1.0;
	}
	//
	//  Apply asymptotic formula if argument is B or greater.
	//
	y = 1.0 / z / z;

	value += 0.5 * y + (1.0 + y * (b2 + y * (b4 + y * (b6 + y * b8)))) / z;

	return value;
}



//****************************************************************************80
//
//  Purpose:
//
//    DIGAMA calculates DIGAMMA ( X ) = d ( LOG ( GAMMA ( X ) ) ) / dX
//
//  Modified:
//
//    18 January 2008
//
//  Author:
//
//    Original FORTRAN77 version by Jose Bernardo.
//    C++ version by John Burkardt.
//
//  Reference:
//
//    Jose Bernardo,
//    Algorithm AS 103:
//    Psi ( Digamma ) Function,
//    Applied Statistics,
//    Volume 25, Number 3, 1976, pages 315-317.
//
//  Parameters:
//
//    Input, double X, the argument of the digamma function.
//    0 < X.
//
//    Output, int *IFAULT, error flag.
//    0, no error.
//    1, X <= 0.
//
//    Output, double DIGAMA, the value of the digamma function at X.
//
double DirichletFit::digamma(double x) {

	double c = 8.5;
	double d1 = -0.5772156649;
	double r;
	double s = 0.00001;
	double s3 = 0.08333333333;
	double s4 = 0.0083333333333;
	double s5 = 0.003968253968;
	double value = 0.0;
	double y = x;
	//
	//  Check the input.
	//
	if (x <= 0.0) return 0.0;

	//
	//  Use approximation if argument <= S.
	//
	if (y <= s) {
		value = d1 - 1.0 / y;
		return value;
	}
	//
	//  Reduce to DIGAMA(X + N) where (X + N) >= C.
	//
	while (y < c) {
		value -= 1.0 / y;
		y += 1.0;
	}
	//
	//  Use Stirling's (actually de Moivre's) expansion if argument > C.
	//
	r = 1.0 / y;
	value += log(y) - 0.5 * r;
	r *= r;
	value -= r * (s3 - r * (s4 - r * s5));

	return value;
}

